# pip install langchain-chroma_本地与云端 langchain llama-index sentence-transformers

import os
import uuid
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.schema import Document

# ========================
# 1️⃣ 基本路径配置
# ========================
base_dir = os.path.dirname(os.path.abspath(__file__))
doc_file_path = os.path.join(base_dir, "data", "manager_learn.txt")

if not os.path.exists(doc_file_path):
    raise FileNotFoundError(f"{doc_file_path} 文件不存在，请检查路径！")

# ========================
# 2️⃣ 加载文本
# ========================
loader = TextLoader(doc_file_path, encoding="utf-8")
documents = loader.load()

# ========================
# 3️⃣ 文本切分
# ========================
text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=20)
split_docs = text_splitter.split_documents(documents)

# ========================
# 4️⃣ 给每条文档加 id
# ========================
split_docs_with_id = [
    Document(page_content=doc.page_content, metadata={"id": str(uuid.uuid4())})
    for doc in split_docs
]

# ========================
# 5️⃣ 初始化 HuggingFace 本地嵌入模型
# ========================
embedding_model = HuggingFaceEmbeddings(
    model_name="D:\\models\\models\\sentence-transformers\\paraphrase-multilingual-MiniLM-L12-v2"
)

# ========================
# 6️⃣ 创建 Chroma 向量库
# ========================
db = Chroma.from_documents(
    documents=split_docs_with_id,
    embedding=embedding_model,
    persist_directory=os.path.join(base_dir, "chroma_db")  # 本地持久化
)
db.persist()

# ========================
# 7️⃣ 相似度搜索
# ========================
query = "贝尔宾的团队角色"
query_docs = db.similarity_search(query, k=5)
print("最相似文档:", query_docs[0].page_content)

# ========================
# 8️⃣ 新增文档
# ========================
new_text = "这是一个新加入的管理学文档内容，用于测试向量库新增。"
new_doc_id = str(uuid.uuid4())
db.add_texts([new_text], metadatas=[{"id": new_doc_id}])
db.persist()
print("新增文档完成！")

# ========================
# 9️⃣ 删除文档（按 id）
# ========================
delete_id = split_docs_with_id[0].metadata["id"]  # 假设删除第一条
db.delete(ids=[delete_id])
db.persist()
print(f"删除文档 id={delete_id} 完成！")

# ========================
# 🔟 更新文档（先删再加）
# ========================
update_id = split_docs_with_id[1].metadata["id"]  # 假设更新第二条
updated_text = "这是更新后的管理学文档内容，替换原文档。"

# 先删除
db.delete(ids=[update_id])
# 再新增，保持 id 不变
db.add_texts([updated_text], metadatas=[{"id": update_id}])
db.persist()
print(f"更新文档 id={update_id} 完成！")
